#!/usr/bin/python
import os
import sys
import numpy as np

# parameters.
if len(sys.argv) < 4:
	print "not enough args"
	sys.exit()
	
pileup_in_dir = sys.argv[1]
pileup_out_dir = sys.argv[2]
contig_size_file = sys.argv[3]

# Read in contig sizes.
print "reading contig sizes."
fin = open(contig_size_file, "rb")
lines = fin.readlines()
fin.close()

tmp = lines[0].split()
try:
	int(tmp[0])
	num = 0
	txt = 1
except:
	num = 1
	txt = 0
	
ctg_sizes = {}
for line in lines:
	tmp = line.strip().split()
	ctg_sizes[tmp[txt]] = int(tmp[num])
del lines

# Get list of input files.
pileup_in_files = os.listdir(pileup_in_dir)
sz_inputs = len(pileup_in_files)

# Build numpy arrays for each contig.
ctg_profiles = {}

# debug variable to detect most hit contigs.
debug = {}

# read in data.
file_idx = {}
fidx = 0
for pileup_in_file in pileup_in_files:
	# Build index of files.
	fpath = "%s/%s" % (pileup_in_dir, pileup_in_file)
	file_idx[fidx] = pileup_in_file
	
	# Open file for reading.	
	fin = open(fpath, "rb")
	
	# Loop over file.
	print "reading %s" % pileup_in_file
	for line in fin:
		# tokenize.
		tmp = line.split()
		ctg = tmp[0]
		idx = int(tmp[1]) - 1
		cov = float(tmp[2])
		
		# check existance.
		if ctg not in ctg_profiles:
			# make appropriate number of arrays.
			tmp1 = []
			for i in range(sz_inputs):
				tmp1.append( np.zeros(ctg_sizes[ctg], dtype=np.float) )
			ctg_profiles[ctg] = tmp1
			
		# debug.
		if ctg not in debug:
			debug[ctg] = 0
		debug[ctg] += 1
			
		# add coverage.
		ctg_profiles[ctg][fidx][idx] = cov
		#print ctg, fidx, idx, ctg_profiles[ctg][fidx][idx]
		
	# Close file for reading.
	fin.close()
	
	# update index.
	fidx += 1
	
# print out profile.
print "printing output."
for ctg in ctg_profiles:
	
	# open output file.
	fout = open("%s/%s.pileup" % (pileup_out_dir, ctg), "wb")
	
	# print header.
	tmp = "#" + file_idx[0]
	for i in range(1,sz_inputs):
		tmp = "%s\t%s" % (tmp, file_idx[i])
	fout.write("%s\n" % tmp)
	
	# loop over each position in contig.
	for i in range(ctg_sizes[ctg]):
		
		# loop over each source.
		tmp = "%f" % ctg_profiles[ctg][0][i]
		for j in range(1, sz_inputs):
			tmp = "%s\t%f" % (tmp, ctg_profiles[ctg][j][i])
		
		# write line to output.
		fout.write("%s\n" % tmp)
	fout.close()
